#ifndef KDTREE_H
#define KDTREE_H

#include <algorithm>
#include <cstdio>
#include <queue>
#include <vector>

#include "base.h"
#include "geo3d.h"

// KD 树，由于网格性质，只在 x,y 方向切分
// 类型 T 需要实现 GetCenter, GetMinBrd, GetMaxBrd 方法
template <typename T>
class KDTree {
public:
    struct Node {
        // 是否为叶子节点
        bool is_leaf;
        // 左右孩子，0为空
        int ch[2];
        // 叶子节点对应的数据
        T* data; 
        // 节点范围
        Point3d min_brd, max_brd;
        Node() {
            is_leaf = false;
            ch[0] = ch[1] = 0;
            data = NULL;
            min_brd = Point3d(1e9, 1e9, 1e9);
            max_brd = Point3d(-1e9, -1e9, -1e9);
        }
    };
    int root;
    std::vector<Node> nodes;
    int stack_top;
    int *node_stack;
    // 建树
    void BuildTree(T* data, int n_data);
    // 测试与直线相交，返回所有潜在相交
    std::vector<T*> IntersectLine(Line line);
private: 
    int BuildTree(T* data, int l, int r, int* index, int depth);
};

template <typename T>
int KDTree<T>::BuildTree(T* data, int l, int r, int* index, int depth) {
    // printf("build tree %d %d %d\n", l, r, depth); fflush(stdout);
    int cur = nodes.size();
    nodes.push_back(Node());
    Node& node = nodes[cur];
    node.is_leaf = false;
    node.min_brd = Point3d(1e9, 1e9, 1e9);
    node.max_brd = Point3d(-1e9, -1e9, -1e9);
    for(int i = l; i < r; i++) {
        int idx = index[i];
        Point3d min_brd = data[idx].GetMinBrd();
        Point3d max_brd = data[idx].GetMaxBrd();
        node.min_brd.x = std::min(node.min_brd.x, min_brd.x);
        node.min_brd.y = std::min(node.min_brd.y, min_brd.y);
        node.min_brd.z = std::min(node.min_brd.z, min_brd.z);
        node.max_brd.x = std::max(node.max_brd.x, max_brd.x);
        node.max_brd.y = std::max(node.max_brd.y, max_brd.y);
        node.max_brd.z = std::max(node.max_brd.z, max_brd.z);
    }
    if(r - l == 1) {
        node.is_leaf = true;
        node.data = &data[index[l]];
        return cur;
    }
    
    double x_mid = (node.min_brd.x + node.max_brd.x) / 2;
    double y_mid = (node.min_brd.y + node.max_brd.y) / 2;
    auto nodeType = [&](int idx) {
        if(depth % 2 == 0) {
            return data[idx].GetCenter().x < x_mid;
        } else {
            return data[idx].GetCenter().y < y_mid;
        }
    };
    // sort index[l...r] according to nodeType
    std::sort(index + l, index + r, [&](int a, int b) {
        return nodeType(a) < nodeType(b);
    });
    // printf("finish sort tree %d %d %d\n", l, r, depth); fflush(stdout);
    int mid = (l + r) >> 1;
    nodes[cur].ch[0] = BuildTree(data, l, mid, index, depth + 1);
    // printf("finish left tree %d %d %d\n", l, r, depth); fflush(stdout);
    nodes[cur].ch[1] = BuildTree(data, mid, r, index, depth + 1);
    // printf("finish right tree %d %d %d\n", l, r, depth); fflush(stdout);
    return cur;
}

template <typename T>
void KDTree<T>::BuildTree(T* data, int n_data) {
    int* index = new int[n_data];
    node_stack = new int[n_data * 2];
    for(int i = 0; i < n_data; i++) {
        index[i] = i;
    }
   nodes.reserve(n_data * 2);
    // nodes[0] as dummy node
    nodes.push_back(Node());

    root = BuildTree(data, 0, n_data, index, 0);
    delete[] index;
}

template <typename T>
std::vector<T*> KDTree<T>::IntersectLine(Line line) {
    std::vector<T*> ret;
    // std::queue<int> q;
    // q.push(root);
    stack_top = 0;
    node_stack[stack_top++] = root;
    while(stack_top) {
        // int cur = q.front();
        // q.pop();
        int cur = node_stack[--stack_top];
        Node& node = nodes[cur];
        if(geo3d::lineBoxIntersection(line, node.min_brd, node.max_brd)) {
            if(node.is_leaf) {
                ret.push_back(node.data);
            } else {
                if(node.ch[0])
                    node_stack[stack_top++] = node.ch[0];
                    // q.push(node.ch[0]);
                if(node.ch[1])
                    node_stack[stack_top++] = node.ch[1];
                    // q.push(node.ch[1]);
            }
        }
    }
    return ret;
}

#endif